from tensorflow.examples.tutorials.mnist import input_data
import tensorflow as tf
from alexnet import model

mnist = input_data.read_data_sets("data_set/data", one_hot=True)

learning_rate = 0.0001
training_iters = 200000
batch_sizes = 128
display_step = 10

n_input = 784
n_classes = 10

x = tf.placeholder(tf.float32, [None, n_input])
y = tf.placeholder(tf.float32, [None, n_classes])
keep_prob = tf.placeholder(tf.float32)

pred = model.alexnet(x)
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost)

correct_pred = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

init = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init)
    step = 1
    while step * batch_sizes < training_iters:
        batch_x, batch_y = mnist.train.next_batch(batch_sizes)
        sess.run(optimizer, feed_dict={x: batch_x, y: batch_y})
        if step % display_step == 0:
            loss, acc = sess.run([cost, accuracy], feed_dict={x: batch_x, y: batch_y})
            print("Iter :" + str(step * batch_sizes) + ", Minibatch loss = "
                  + str(loss) + ", Training_accuracy = " + str(acc))

        step += 1

    print("finished")
    print("Test Accuracy:", sess.run(accuracy, feed_dict={x: mnist.test.images[:256], y: mnist.test.labels[:256]}))
